/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2021-2021.
 * Description: Support for dynamic loading of kernel symbols.
 * Author: hanzongcheng
 * Create: 2021-06-15
 */
#include <linux/atomic.h>
#include <linux/stop_machine.h>
#include <linux/init.h>
#ifdef CONFIG_KPROBES
#include <linux/kprobes.h>
#endif
#ifdef CONFIG_FUNCTION_ERROR_INJECTION
#include <linux/error-injection.h>
#endif
#include <linux/rtos/extract_kallsyms.h>
#ifdef CONFIG_RTOS_HKRR
#include <linux/slab.h>

#include "../../hkrr_utils.h"
#endif

static atomic_t initialized = ATOMIC_INIT(0);

#ifdef CONFIG_RTOS_HKRR
static u8 *kas_names;
struct hkrr_kastab *kastab;

static int alloc_hkrr_mem(struct rtos_kallsyms_struct *symbuf)
{
	int kas_names_size;

	if (*(symbuf->num_syms) <= 0)
		return -EINVAL;

	kastab = kcalloc(*(symbuf->num_syms), sizeof(struct hkrr_kastab), GFP_KERNEL);
	if (!kastab)
		return -ENOMEM;

	kas_names_size = (const u8 *)(symbuf->markers) - symbuf->names;
	kas_names = kcalloc(kas_names_size, sizeof(u8), GFP_KERNEL);
	if (!kas_names) {
		kfree(kastab);
		kastab = NULL;
		return -ENOMEM;
	}

	return 0;
}

static void cleanup_hkrr_mem(void)
{
	kfree(kastab);
	kfree(kas_names);
	kastab = NULL;
	kas_names = NULL;
}
#endif

static int update_kallsyms_buf(void *data)
{
	struct rtos_kallsyms_struct *symbuf = data;
#if !defined(CONFIG_X86) && !defined(CONFIG_ARM64)
#if defined(CONFIG_STRICT_KERNEL_RWX) || defined(CONFIG_RTOS_PPC_TEXT_PROTECT)
	rtos_mark_rodata_rw();
#endif
#endif
	kallsyms_names = symbuf->names;
	kallsyms_token_table = symbuf->token_table;
	kallsyms_token_index = symbuf->token_index;
	kallsyms_markers = symbuf->markers;
	kallsyms_offsets = symbuf->offsets;
	kallsyms_num_syms = *(symbuf->num_syms);
#ifdef CONFIG_RTOS_HKRR
	hkrr_kallsyms_data_setup(kas_names, kastab);
#endif
#if !defined(CONFIG_X86) && !defined(CONFIG_ARM64)
#if defined(CONFIG_STRICT_KERNEL_RWX) || defined(CONFIG_RTOS_PPC_TEXT_PROTECT)
	rtos_mark_rodata_ro();
#endif
#endif
	return 0;
}

int rtos_kallsyms_init(struct rtos_kallsyms_struct *symbuf)
{
	int ret = 0;

	if (!atomic_add_unless(&initialized, 1, 1))
		return -EBUSY;

#ifdef CONFIG_RTOS_HKRR
	ret = alloc_hkrr_mem(symbuf);
	if (ret)
		return ret;
#endif

	/*
	 * Before updating symbuf, the rodata needs to be set to RW.
	 * Make sure the rodata is updated with preemption disabled for
	 * all architecures except x86.
	 * On x86 architecture, rtos_mark_rodata_rw() and rtos_mark_rodata_ro()
	 * might sleep, so we do not modify rodata in the stop_machine process.
	 */
#if (defined(CONFIG_X86) || defined(CONFIG_ARM64)) && defined(CONFIG_STRICT_KERNEL_RWX)
	rtos_mark_rodata_rw();
#endif
	stop_machine(update_kallsyms_buf, symbuf, NULL);
#if (defined(CONFIG_X86) || defined(CONFIG_ARM64)) && defined(CONFIG_STRICT_KERNEL_RWX)
	rtos_mark_rodata_ro();
#endif
#ifdef CONFIG_KPROBES
	rtos_populate_kprobe_blacklist();
#endif
#ifdef CONFIG_FUNCTION_ERROR_INJECTION
	rtos_populate_kernel_ei_list();
#endif
#ifdef CONFIG_RTOS_HKRR
	cleanup_hkrr_mem();
#endif
	return ret;
}
EXPORT_SYMBOL_NS(rtos_kallsyms_init, RTOS_EXTRACT_KALLSYMS);
